/*
 * Copyright 2013-2019 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.springframework.cloud.gateway.filter.factory;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

import cn.hutool.core.lang.UUID;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.tsfyun.scm.gateway.config.properties.NoticeProperties;
import com.tsfyun.scm.gateway.constant.LimitConstant;
import com.tsfyun.scm.gateway.support.dd.DingTalkNoticeUtil;
import com.tsfyun.scm.gateway.util.FrequencyUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;

/**
 * User Request Rate Limiter filter. See https://stripe.com/blog/rate-limiters and
 */
@RefreshScope
@ConfigurationProperties("spring.cloud.gateway.filter.request-rate-limiter")
@Slf4j
public class RequestRateLimiterGatewayFilterFactory extends
        AbstractGatewayFilterFactory<RequestRateLimiterGatewayFilterFactory.Config> {

    @Autowired
    private NoticeProperties noticeProperties;

    @Value("${gateway.manage.salt:asdfg12345}")
    private String salt;

    @Value("${gateway.frequency.timeout:600}")
    private Integer frequencyTimeout;

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    /**
     * Key-Resolver key.
     */
    public static final String KEY_RESOLVER_KEY = "keyResolver";

    private static final String EMPTY_KEY = "____EMPTY_KEY__";

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    /**
     * Switch to deny requests if the Key Resolver returns an empty key, defaults to true.
     */
    private boolean denyEmptyKey = true;

    /** HttpStatus to return when denyEmptyKey is true, defaults to FORBIDDEN. */
    private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();

    public RequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter,
                                                  KeyResolver defaultKeyResolver) {
        super(Config.class);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
    }

    public KeyResolver getDefaultKeyResolver() {
        return defaultKeyResolver;
    }

    public RateLimiter getDefaultRateLimiter() {
        return defaultRateLimiter;
    }

    public boolean isDenyEmptyKey() {
        return denyEmptyKey;
    }

    public void setDenyEmptyKey(boolean denyEmptyKey) {
        this.denyEmptyKey = denyEmptyKey;
    }

    public String getEmptyKeyStatusCode() {
        return emptyKeyStatusCode;
    }

    public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
        this.emptyKeyStatusCode = emptyKeyStatusCode;
    }

    @SuppressWarnings("unchecked")
    @Override
    public GatewayFilter apply(Config config) {
        KeyResolver resolver = getOrDefault(config.keyResolver, defaultKeyResolver);
        RateLimiter<Object> limiter = getOrDefault(config.rateLimiter,
                defaultRateLimiter);
        boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder
                .parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));

        return (exchange, chain) -> {
            Route route = exchange
                    .getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);

            return resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
                if (EMPTY_KEY.equals(key)) {
                    if (denyEmpty) {
                        setResponseStatus(exchange, emptyKeyStatus);
                        return exchange.getResponse().setComplete();
                    }
                    return chain.filter(exchange);
                }
                return limiter.isAllowed(route.getId(), key).flatMap(response -> {

                    for (Map.Entry<String, String> header : response.getHeaders()
                            .entrySet()) {
                        exchange.getResponse().getHeaders().add(header.getKey(),
                                header.getValue());
                    }

                    if (response.isAllowed()) {
                        return chain.filter(exchange);
                    }
                    log.info("进入网关限流拦截，token【{}】,请求路径【{}】",exchange.getRequest().getHeaders().getFirst("token"),exchange.getRequest().getURI());
                    //发送钉钉通知
                    List<String> dingdingUrls = noticeProperties.getDingdingUrl();
                    DingTalkNoticeUtil.send2DingDing(exchange.getRequest(),dingdingUrls,salt);
                    ServerHttpResponse serverHttpResponse = exchange.getResponse();
                    serverHttpResponse.setStatusCode(HttpStatus.OK);

                    String limitRequestId = exchange.getRequest().getHeaders().getFirst("limitRequestId");
                    if(StringUtils.isEmpty(limitRequestId) || Objects.equals("null", StrUtil.nullToEmpty(limitRequestId).toLowerCase())) {
                        limitRequestId = UUID.fastUUID().toString();
                    }
                    stringRedisTemplate.opsForValue().set(LimitConstant.FREQUENCY_LIMIT + limitRequestId,limitRequestId,frequencyTimeout, TimeUnit.SECONDS);
                    JSONObject resultJson = FrequencyUtil.wrapResult(limitRequestId);
                    byte[] bytes = resultJson.toJSONString().getBytes(StandardCharsets.UTF_8);
                    exchange.getResponse().getHeaders().add("Content-Type","application/json;charset=UTF-8");
                    DataBuffer buffer = exchange.getResponse().bufferFactory().wrap(bytes);
                    return exchange.getResponse().writeWith(Flux.just(buffer));
                });
            });
        };
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }

    public static class Config {

        private KeyResolver keyResolver;

        private RateLimiter rateLimiter;

        private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;

        private Boolean denyEmptyKey;

        private String emptyKeyStatus;

        public KeyResolver getKeyResolver() {
            return keyResolver;
        }

        public Config setKeyResolver(KeyResolver keyResolver) {
            this.keyResolver = keyResolver;
            return this;
        }

        public RateLimiter getRateLimiter() {
            return rateLimiter;
        }

        public Config setRateLimiter(RateLimiter rateLimiter) {
            this.rateLimiter = rateLimiter;
            return this;
        }

        public HttpStatus getStatusCode() {
            return statusCode;
        }

        public Config setStatusCode(HttpStatus statusCode) {
            this.statusCode = statusCode;
            return this;
        }

        public Boolean getDenyEmptyKey() {
            return denyEmptyKey;
        }

        public Config setDenyEmptyKey(Boolean denyEmptyKey) {
            this.denyEmptyKey = denyEmptyKey;
            return this;
        }

        public String getEmptyKeyStatus() {
            return emptyKeyStatus;
        }

        public Config setEmptyKeyStatus(String emptyKeyStatus) {
            this.emptyKeyStatus = emptyKeyStatus;
            return this;
        }

    }

}
